Skip to content

[None][fix] harden Python native KV transfer ownership#16396

Draft
chienchunhung wants to merge 5 commits into
NVIDIA:mainfrom
chienchunhung:codex/python-native-kv-transfer-ownership
Draft

[None][fix] harden Python native KV transfer ownership#16396
chienchunhung wants to merge 5 commits into
NVIDIA:mainfrom
chienchunhung:codex/python-native-kv-transfer-ownership

Conversation

@chienchunhung

@chienchunhung chienchunhung commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Important

Integration draft — not ready for code review as one PR. The current branch combines several ownership layers and will be decomposed into review-sized production+test slices. Please use #16347 for the design contract; do not review this aggregate diff yet.

Summary

  • implement the initial Python-runtime containment slice of the ownership contract defined in [None][docs] define Python native KV transfer ownership contract #16347
  • add registry-owned receive lifecycle state with exact writer-cohort reconciliation, publication/cancellation gating, idempotent result routing, and separate logical/physical drain state
  • retain sender/receiver request roots, NIXL operation state, descriptor state, bounce leases, registrations, and mappings until terminal evidence permits release
  • make cancellation and executor/resource-manager teardown fail closed while a Python-native transfer owner is still live
  • harden V1 block-pin retirement with transactional unpinBlocksById validation
  • add focused lifecycle, race, cancellation, replay, bounce, and executor teardown coverage

Relationship to #16347

This PR is the first implementation follow-up to #16347 and is logically stacked on it. The #16347 head branch lives in the fork, so this PR targets main; until #16347 merges, GitHub will also show its three design-document changes in this PR.

#16347 This PR
Deliverable Design and phased implementation plan only Runtime containment code, tests, and implementation-status documentation
Target contract Complete allocator- and protocol-backed ownership lifecycle Initial fail-closed containment using the existing #15618 wire format
Ownership mechanism Generation-safe endpoint records plus allocator-issued KV/bounce leases Registry records, exact writer ledgers, strong request/manager roots, V1 block pins, and bounce-slot settlement
Retirement Per-operation quiescence or negotiated/global fence Known terminal evidence; ambiguous work retains the owner and mappings
Remaining work Defines the full target Allocation-generation KV leases, transfer generations/endpoint epochs, bounded replay, capability negotiation, terminal acknowledgements, and definitive remote/global quiescence

Applicability

  • applies only when transceiver_runtime="PYTHON" selects Python-native KvCacheTransceiverV2
  • covers direct transfers and eligible bounce transfers; it is not gated by kv_cache_bounce_size_mb or a second ownership flag
  • leaves the default C++ CacheTransceiver, its data path, and its wire protocol unchanged
  • keeps generation-first attention-DP broadcast on the compatibility path and outside the ownership rollout
  • disables receive bounce for PP fan-in until byte-accurate per-writer extents can be authorized before address publication
  • preserves the [TRTLLM-13614][feat] Disaggregated KV-cache bounce transfer #15618 wire format, so mixed-version/protocol hardening is not claimed here

Safety and operational impact

The normal path gains lifecycle bookkeeping and exact-set validation but no additional device arena beyond an already-active bounce configuration. Under ambiguous transport failure, the implementation deliberately retains requests, KV allocations/pins, bounce slots, and transport mappings instead of risking reuse or use-after-free. That converts a potential memory-safety failure into visible capacity pressure or a teardown veto; bounded recovery requires the follow-up lease/protocol work described in #16347.

Validation

  • iterative trtllm-review self-review through a final clean pass with no P0-P2 findings
  • full changed-file pre-commit run --from-ref FETCH_HEAD --to-ref HEAD
  • git diff --check
  • Python bytecode compilation for all changed Python files
  • focused dependency-isolated lifecycle tests: 29 passed
  • focused executor ownership/scope checks: 7 passed
  • focused post-rebase C++/Python cancellation interaction checks: 8 passed
  • git range-diff confirmed all three patches were unchanged when rebased onto current main (d17ba0e4d)
  • test-list stage mapping validation

Not run locally: the full pytest selection (the environment lacks transformers), compiled C++ unit tests, or GPU/NIXL/VMM/PP/ADP/multirank runtime tests. Those remain required before this draft is ready for merge.

@chienchunhung
chienchunhung force-pushed the codex/python-native-kv-transfer-ownership branch 2 times, most recently from 3b5f743 to 3862973 Compare July 15, 2026 02:48

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59340 [ run ] triggered by Bot. Commit: 3862973 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59340 [ run ] completed with state SUCCESS. Commit: 3862973
/LLM/main/L0_MergeRequest_PR pipeline #47820 completed with status: 'ABORTED'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
@chienchunhung
chienchunhung force-pushed the codex/python-native-kv-transfer-ownership branch from 3862973 to 6888a21 Compare July 22, 2026 05:15
@chienchunhung

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60892 [ run ] triggered by Bot. Commit: 6888a21 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #60892 [ run ] completed with state FAILURE. Commit: 6888a21
/LLM/main/L0_MergeRequest_PR pipeline #49159 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
@chienchunhung

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast --stage-list "A10-PyTorch-2"

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61124 [ run ] triggered by Bot. Commit: 7b4bc96 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61124 [ run ] completed with state SUCCESS. Commit: 7b4bc96
/LLM/main/L0_MergeRequest_PR pipeline #49379 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants